import os
os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"
import random
from torchvision import datasets, transforms

import numpy as np
import torch
import os
import matplotlib.pyplot as plt
#设置随机数种子，保证结果的可复现
def setup_seed(seed):
    np.random.seed(seed)
    random.seed(seed)
    os.environ['PYTHONHASHSEED'] = str(seed)
    torch.manual_seed(seed)

#设置随机数种子
setup_seed(26)

#加载数据,ToTensor()将数据转换为tensor
train_data = datasets.MNIST(root='./data', train=True,download=True, transform=transforms.ToTensor())
test_data = datasets.MNIST(root='./data',train=False,download=True, transform=transforms.ToTensor())

#数据切分
train_loader = torch.utils.data.DataLoader(train_data, batch_size=64, shuffle=True)
test_loader = torch.utils.data.DataLoader(test_data, batch_size=64, shuffle=True)

#显示6个图片
eg = enumerate(train_loader)
batch,(img,labels) = next(eg)
for i in range(6):
    plt.subplot(2,3,i+1)
    plt.imshow(img[i][0],cmap= 'gray',interpolation='none')
    plt.title(f'truth:{labels[i]}')
plt.show()